package com.love.tree;

/**
 * 前缀树的实现
 */
public class Trie {

    public TrieNode root ;

    class TrieNode{

        /**
         * 表示从某个路径开始 [某个字符] 出现了几次
         */
        public int pass;

        /**
         * 表示某个单词出现了几次
         */
        public int end;

        public TrieNode[] nexts;

        public TrieNode(){
            pass = 0;
            end = 0;
            nexts = new TrieNode[26];
        }


    }

    public Trie(){
        root = new TrieNode();
    }

    public void insert(String word){
        TrieNode node = root;
        node.pass++;
        for (int i = 0,path; i < word.length(); i++) {
            path = word.charAt(i) - 'a';
            if(node.nexts[path] == null){
                node.nexts[path] = new TrieNode();
            }
            node = node.nexts[path];
            node.pass++;
        }
        node.end++;
    }

    /**
     * 查询前缀树里，word单次出现了几次
     * @param word
     * @return
     */
    public int countWordsEqualTo(String word){
        TrieNode node = root;
        for (int i = 0,path; i < word.length(); i++) {
            path = word.charAt(i) - 'a';
            if(node.nexts[path] == null){
                return 0;
            }
            node = node.nexts[path];
        }
        return node.end;
    }

    /**
     * 查询前缀树里，有多少单词以pre做前缀
     * @param pre
     * @return
     */
    public int countWordsStartWith(String pre){
        TrieNode node = root;
        for (int i = 0,path; i < pre.length(); i++) {
            path = pre.charAt(i) - 'a';
            if(node.nexts[path] == null){
                return 0;
            }
            node = node.nexts[path];
        }
        return node.pass;
    }


    /**
     * 如果之前word插入过前缀树，那么此时删除一次
     * 如果之前word没有插入过前缀树，那么什么也不做
     * @param word
     */
    public void erase(String word){
        if(countWordsEqualTo(word) > 0){
            TrieNode node = root;
            node.pass--;
            for (int i = 0,path; i < word.length(); i++) {
                path = word.charAt(i) - 'a';
                if(--node.nexts[path].pass == 0){
                    node.nexts[path] = null;
                    return;
                }
                node = node.nexts[path];
            }
            node.end--;
        }
    }


}
